import math
import torch
import torch.nn as nn

def get_activation(name="silu", inplace=True):
    if name == "silu":
        module = nn.SiLU(inplace=inplace)
    elif name == "relu":
        module = nn.ReLU(inplace=inplace)
    elif name == "lrelu":
        module = nn.LeakyReLU(0.1, inplace=inplace)
    elif name == "hardswish":
        module = nn.Hardswish(inplace=inplace)
    else:
        raise AttributeError("Unsupported act type: {}".format(name))
    return module

class BaseConv(nn.Module):
    def __init__(
        self, in_channels, out_channels, ksize, stride, groups=1, bias=False, act="silu"
    ):
        super().__init__()
        # same padding
        pad = (ksize - 1) // 2
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=ksize,
            stride=stride,
            padding=pad,
            groups=groups,
            bias=bias,
        )
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = get_activation(act, inplace=True)

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))
    
class DWConv(nn.Module):
    def __init__(self, in_channels, out_channels, ksize, stride=1, act="silu"):
        super().__init__()
        self.dconv = BaseConv(
            in_channels,
            in_channels,
            ksize=ksize,
            stride=stride,
            groups=in_channels,
            act=act,
        )
        self.pconv = BaseConv(
            in_channels, out_channels, ksize=1, stride=1, groups=1, act=act
        )

    def forward(self, x):
        x = self.dconv(x)
        return self.pconv(x)
    
class CA_Block(nn.Module):
    def __init__(self, in_channels, out_channels, ratio=0.5, act="hardswish"):
        super().__init__()
        hidden_channels = int(in_channels * ratio)
        self.pool_w, self.pool_h = nn.AdaptiveAvgPool2d((1, None)), nn.AdaptiveAvgPool2d((None, 1))
        self.conv1 = nn.Conv2d(in_channels, hidden_channels, kernel_size=1, stride=1)
        self.bn = nn.BatchNorm2d(hidden_channels)
        self.act = get_activation(act, inplace=True)
        self.conv2_h = nn.Conv2d(hidden_channels, out_channels, kernel_size=1, stride=1)
        self.conv2_w = nn.Conv2d(hidden_channels, out_channels, kernel_size=1, stride=1)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        h = x.shape[2]
        w = x.shape[3]
        h_avg_output = self.pool_h(x).permute(0, 1, 3, 2)
        w_ave_output = self.pool_w(x)
        conv_output1 = self.conv1(torch.concat([h_avg_output, w_ave_output], dim=3))
        F_h, F_w = self.act(self.bn(conv_output1)).split([h, w], 3)
        h_output = self.sigmoid(self.conv2_h(F_h.permute(0, 1, 3, 2)))
        w_output = self.sigmoid(self.conv2_w(F_w))

        return x * h_output.expand_as(x) * w_output.expand_as(x)